import socket
import sys
from enum import Enum
from typing import List

import psutil
from PyQt5 import QtWidgets
from PyQt5.QtCore import pyqtSignal, pyqtSlot
from PyQt5.QtNetwork import QTcpServer, QHostAddress, QTcpSocket
from PyQt5.QtWidgets import QApplication, QMessageBox
from pydantic import BaseModel


class ServeType(Enum):
    server_tcp = "TCP服务器"
    server_udp = "UDP服务器"
    client_tcp = "TCP客户端"
    client_udp = "UDP客户端"


# 客户选消息模型
class ClientMsgModel(BaseModel):
    client_ip: str
    client_port: int
    msg: str


class UiTcpUdp(QtWidgets.QWidget):
    # signal_write_msg = pyqtSignal(str)

    def __init__(self, stop_thread=None, num=1):
        """
        初始化，定义变量
        @param stop_thread: StopThreading类创建的对象
        @param num: 多开的窗口数，第一个开的窗口未1
        """
        super().__init__()
        self.server_tcp = None
        self.another = None
        self.st = None
        self.network_interfaces: dict = {}
        self.num = num
        self.stop_thread = stop_thread
        self.wins = dict()

        self.resize(640, 480)
        self.setWindowTitle("TCP/UDP网络测试工具-窗口%s" % self.num)
        self.setAcceptDrops(False)
        # self.setSizeGripEnabled(False)

        # 组件
        self.label_interface: QtWidgets.QLabel = QtWidgets.QLabel()  # 选择网卡
        self.combox_interface: QtWidgets.QComboBox = QtWidgets.QComboBox()  # 选择网卡
        self.edit_ip_local: QtWidgets.QLineEdit = QtWidgets.QLineEdit()  # 网卡对应IP
        self.button_get_ip: QtWidgets.QPushButton = (
            QtWidgets.QPushButton()
        )  # 重新获取网卡及IP对应关系

        self.edit_port: QtWidgets.QLineEdit = QtWidgets.QLineEdit()  # 端口号
        self.label_port: QtWidgets.QLabel = QtWidgets.QLabel()

        self.edit_ip_send: QtWidgets.QLineEdit = QtWidgets.QLineEdit()  # 目标IP
        self.label_ip_send: QtWidgets.QLabel = QtWidgets.QLabel()

        self.edit_send: QtWidgets.QTextEdit = QtWidgets.QTextEdit()  # 发送数据区域
        self.label_send: QtWidgets.QLabel = QtWidgets.QLabel()

        self.button_multipleOpening: QtWidgets.QPushButton = (
            QtWidgets.QPushButton()
        )  # 多开
        self.button_connect: QtWidgets.QPushButton = QtWidgets.QPushButton()  # 连接
        self.button_disconnect: QtWidgets.QPushButton = QtWidgets.QPushButton()  # 断开连接
        self.combox_server_type: QtWidgets.QComboBox = QtWidgets.QComboBox()  # 服务类型

        self.button_send: QtWidgets.QPushButton = QtWidgets.QPushButton()  # 发送按钮
        self.button_exit: QtWidgets.QPushButton = QtWidgets.QPushButton()  # 退出

        self.label_recv_all: QtWidgets.QLabel = QtWidgets.QLabel()  # 接收区域

        self.label_recv: QtWidgets.QLabel = QtWidgets.QLabel()  # 接收区域
        self.browser_recv: QtWidgets.QTextBrowser = QtWidgets.QTextBrowser()

        self.label_client: QtWidgets.QLabel = QtWidgets.QLabel()  # 客户端
        self.list_client: QtWidgets.QListWidget = QtWidgets.QListWidget()  # 客户端列表

        self.button_clean: QtWidgets.QPushButton = QtWidgets.QPushButton()  # 清空接收区域

        # 整体
        self.box_all: QtWidgets.QHBoxLayout = QtWidgets.QHBoxLayout()

        # 左右分列
        self.box_left: QtWidgets.QVBoxLayout = QtWidgets.QVBoxLayout()
        self.box_right: QtWidgets.QVBoxLayout = QtWidgets.QVBoxLayout()

        # 功能分区
        self.box_set: QtWidgets.QVBoxLayout = QtWidgets.QVBoxLayout()
        self.box_send: QtWidgets.QVBoxLayout = QtWidgets.QVBoxLayout()
        self.box_exit: QtWidgets.QHBoxLayout = QtWidgets.QHBoxLayout()
        self.box_client: QtWidgets.QVBoxLayout = QtWidgets.QVBoxLayout()
        self.box_recv: QtWidgets.QHBoxLayout = QtWidgets.QHBoxLayout()
        self.box_recv_msg: QtWidgets.QVBoxLayout = QtWidgets.QVBoxLayout()

        # 排列
        self.box_h_1: QtWidgets.QHBoxLayout = QtWidgets.QHBoxLayout()
        self.box_h_2: QtWidgets.QHBoxLayout = QtWidgets.QHBoxLayout()
        self.box_h_3: QtWidgets.QHBoxLayout = QtWidgets.QHBoxLayout()
        self.box_h_4: QtWidgets.QHBoxLayout = QtWidgets.QHBoxLayout()

        # 初始化
        self.click_get_ip()
        self.ui_connect()
        self.ui_translate()
        self.ui_layout()

    def ui_translate(self):
        """
        初始化控件
        @return:
        """
        # 本地
        self.label_interface.setText("选择网卡")
        self.button_get_ip.setText("重新获取")

        self.label_port.setText("端口号")
        self.label_ip_send.setText("目标IP")
        self.label_send.setText("发送数据区域：")
        self.label_recv_all.setText("接收数据区域：")
        self.label_recv.setText("接收数据：")
        self.label_client.setText("客户端列表：")

        self.button_clean.setText("清空接收区域")
        self.button_connect.setText("连接")
        self.button_disconnect.setText("断开连接")
        self.button_multipleOpening.setText("多开")
        self.button_send.setText("发送数据")
        self.button_exit.setText("退出系统")

        self.browser_recv.insertPlainText("这是窗口-%s\n" % self.num)

        for server_type in ServeType:
            self.combox_server_type.addItem(server_type.value)

        # self.list_client.addItem("这是窗口-%s" % self.num)
        # self.list_client.addItem("这是窗口-%s" % self.num)

    def ui_layout(self):
        """
        界面排列
        @return:
        """
        # 左边
        self.box_h_1.addWidget(self.label_interface)
        self.box_h_1.addWidget(self.combox_interface)
        self.box_h_1.addWidget(self.edit_ip_local)
        self.box_h_1.addWidget(self.button_get_ip)

        self.box_h_2.addWidget(self.label_port)
        self.box_h_2.addWidget(self.edit_port)
        self.box_h_2.addWidget(self.button_multipleOpening)

        self.box_h_3.addWidget(self.label_ip_send)
        self.box_h_3.addWidget(self.edit_ip_send)

        self.box_h_4.addWidget(self.combox_server_type)
        self.box_h_4.addWidget(self.button_connect)
        self.box_h_4.addWidget(self.button_disconnect)

        self.box_set.addLayout(self.box_h_1)
        self.box_set.addLayout(self.box_h_2)
        self.box_set.addLayout(self.box_h_3)
        self.box_set.addLayout(self.box_h_4)

        self.box_send.addWidget(self.label_send)
        self.box_send.addWidget(self.edit_send)

        self.box_exit.addWidget(self.button_send)
        self.box_exit.addWidget(self.button_exit)

        self.box_left.addLayout(self.box_set)
        self.box_left.addLayout(self.box_send)
        self.box_left.addLayout(self.box_exit)

        # 右边
        self.box_recv.addWidget(self.label_recv_all)
        self.box_recv.addWidget(self.button_clean)

        self.box_client.addWidget(self.label_client)
        self.box_client.addWidget(self.list_client)

        self.box_recv_msg.addWidget(self.label_recv)
        self.box_recv_msg.addWidget(self.browser_recv)

        self.box_right.addLayout(self.box_recv)
        self.box_right.addLayout(self.box_client)
        self.box_right.addLayout(self.box_recv_msg)

        # 整体
        self.box_all.addLayout(self.box_left)
        self.box_all.addLayout(self.box_right)

        self.setLayout(self.box_all)

    def ui_connect(self):
        """
        槽连接
        @return:
        """
        self.button_get_ip.clicked.connect(self.click_get_ip)  # 获取网卡及IP
        self.combox_interface.currentIndexChanged.connect(  # 联动设置IP
            self.interface_change
        )
        self.button_multipleOpening.clicked.connect(self.another_window)
        self.combox_server_type.currentIndexChanged.connect(self.ui_change)
        self.button_connect.clicked.connect(self.link)
        self.button_disconnect.clicked.connect(self.un_link)

    def ui_click_connect(self):
        """
        点击连接后界面变化
        @return:
        """
        # 所有输入框不能点击或输入
        self.combox_interface.setEnabled(False)
        self.edit_ip_local.setEnabled(False)
        self.edit_port.setEnabled(False)
        self.edit_ip_send.setEnabled(False)
        self.edit_send.setEnabled(False)
        self.button_get_ip.setEnabled(False)
        self.button_multipleOpening.setEnabled(False)
        self.button_connect.setEnabled(False)
        self.button_disconnect.setEnabled(True)

    def ui_click_disconnect(self):
        """
        点击断开连接后界面变化
        @return:
        """
        # 恢复所有输入框
        self.combox_interface.setEnabled(True)
        self.edit_ip_local.setEnabled(True)
        self.edit_port.setEnabled(True)
        self.edit_ip_send.setEnabled(True)
        self.edit_send.setEnabled(True)
        self.button_get_ip.setEnabled(True)
        self.button_multipleOpening.setEnabled(True)
        self.button_connect.setEnabled(True)
        self.button_disconnect.setEnabled(False)

    @pyqtSlot()
    def un_link(self):
        """
        断开连接
        @return:
        """
        match self.combox_server_type.currentText():
            case ServeType.server_tcp.value:
                self.server_tcp.close()
                self.ui_click_disconnect()
            case ServeType.server_udp.value:
                pass
            case ServeType.client_tcp.value:
                pass
            case ServeType.client_udp.value:
                pass
            case _:
                pass

    @pyqtSlot()
    def link(self):
        """
        连接
        @return:
        """
        try:
            self.ui_click_connect()

            if self.edit_port.value is not None and self.edit_port.value is not None:
                match self.combox_server_type.currentText():
                    case ServeType.server_tcp.value:
                        __host = self.edit_ip_local.text()
                        __port = int(self.edit_port.text())

                        print(__host, __port)
                        self.server_tcp = TcpServer()
                        self.server_tcp.listen(QHostAddress(__host), __port)
                        # 端口占用
                        if self.server_tcp.isListening():
                            QMessageBox.information(
                                self, "提示", "TCP服务器启动成功！", QMessageBox.Yes
                            )
                            self.server_tcp.signals_connected.connect(  # 新的客户端连接
                                self.handle_new_client
                            )
                            self.server_tcp.signals_disconnected.connect(  # 客户端断开连接
                                self.handle_disconnect
                            )
                            self.server_tcp.signals_message.connect(self.handle_message)
                        else:
                            QMessageBox.warning(self, "警告", "TCP服务器启动失败！", QMessageBox.Yes)
                            # 恢复所有输入框
                            self.button_connect.setEnabled(True)
                            self.button_disconnect.setEnabled(False)
                            self.ui_click_disconnect()
                    case ServeType.server_udp.value:
                        pass
                    case ServeType.client_tcp.value:
                        pass
                    case ServeType.client_udp.value:
                        pass
                    case _:
                        pass


        except Exception as e:
            QMessageBox.error(self, "错误", "连接失败！\n%s" % e, QMessageBox.Yes)

            # 恢复所有输入框
            self.ui_click_disconnect()
        # 直接关闭程序
        except KeyboardInterrupt:
            self.close()

    @pyqtSlot()
    def handle_new_client(self):
        """
        处理新的客户端连接
        @return:
        """

    @pyqtSlot()
    def handle_disconnect(self):
        """
        处理客户端断开连接
        @return:
        """

    @pyqtSlot()
    def handle_message(self):
        """
        处理消息
        @return:
        """

    @pyqtSlot()
    def ui_change(self):
        """
        选择服务类型
        @return:
        """
        match self.combox_server_type.currentText():
            case ServeType.server_tcp.value | ServeType.server_udp.value:
                self.edit_ip_send.hide()
                self.label_ip_send.hide()
                self.edit_port.show()
                self.label_port.show()
                self.label_port.setText("端口号")
                self.button_connect.setText("启动")
                self.button_disconnect.setText("停止")
                self.button_multipleOpening.hide()
            case ServeType.client_tcp.value | ServeType.client_udp.value:
                self.edit_ip_send.show()
                self.label_ip_send.show()
                self.edit_port.show()
                self.label_port.show()
                self.label_port.setText("目标端口")
                self.button_connect.setText("连接")
                self.button_disconnect.setText("断开连接")
                self.button_multipleOpening.show()
            case _:
                pass

    @pyqtSlot()
    def click_get_ip(self):
        """
        pushbutton_get_ip控件点击触发的槽
        @return:
        """
        self.network_interfaces.clear()
        interfaces = psutil.net_if_addrs()

        for interface_name, interface_addresses in interfaces.items():
            for address in interface_addresses:
                if address.family == socket.AF_INET:
                    self.network_interfaces[interface_name] = address.address

        for interface, _ in self.network_interfaces.items():
            self.combox_interface.addItem(interface)

        print(f"窗口：{self.num}  interFaceInfo: {self.network_interfaces}")

    @pyqtSlot()
    def interface_change(self):
        """
        修改网卡
        combobox控件内容改变时触发的槽
        @return:
        """
        __ip = self.network_interfaces[self.combox_interface.currentText()]
        print(f"IP address: {__ip}")
        self.edit_ip_local.setText(__ip)

    @pyqtSlot()
    def another_window(self):
        """
        新开窗口
        @return:
        """
        QMessageBox.warning(
            self, "TCP/UDP网络测试助手", "已经开启了新的TCP/UDP网络测试助手！", QMessageBox.Yes
        )
        self.num = self.num + 1
        self.st = None
        self.another = UiTcpUdp(self.st, self.num)
        self.another.show()

        self.wins[self.num] = self.another


class TcpServer(QTcpServer):
    """Tcp服务器"""

    signals_connected = pyqtSignal(list)
    signals_disconnected = pyqtSignal(list)
    signals_message = pyqtSignal(ClientMsgModel)

    def __init__(self, parent: object = None):
        super().__init__(parent)
        self.client_socket_list = List[QTcpSocket]

    def incomingConnection(self, socketDescriptor):
        """
        重写incomingConnection方法
        :param socketDescriptor:
        :return:
        """
        client_socket = QTcpSocket()
        client_socket.setSocketDescriptor(socketDescriptor)
        self.client_socket_list.append(client_socket)
        self.signals_connected.emit(
            [
                client_socket.peerAddress().toString(),
                client_socket.peerPort(),
            ]
        )
        client_socket.readyRead.connect(self.handle_message)
        client_socket.disconnected.connect(self.handle_disconnect)

    @pyqtSlot()
    def handle_message(self):
        """
        处理消息
        :return:
        """
        client_socket = self.sender()
        data = client_socket.readAll().data().decode()
        __client_msg = ClientMsgModel(  # 客户端消息
            client_ip=client_socket.peerAddress().toString(),
            client_port=client_socket.peerPort(),
            msg=data,
        )
        self.signals_message.emit(__client_msg)

    @pyqtSlot()
    def handle_disconnect(self):
        """
        处理断开连接
        :return:
        """
        client_socket = self.sender()
        self.signals_disconnected.emit(
            [
                client_socket.peerAddress().toString(),
                client_socket.peerPort(),
            ]
        )
        self.client_socket_list.remove(client_socket)
        client_socket.deleteLater()

    @pyqtSlot()
    def send_message(self, client_msg: ClientMsgModel):
        """
        发送消息
        :param client_msg: 客户端消息
        :return:
        """
        for client_socket in self.client_socket_list:
            if (
                    client_socket.peerAddress().toString() == client_msg.client_ip
                    and client_socket.peerPort() == client_msg.client_port
            ):
                client_socket.write(client_msg.msg.encode())

    @pyqtSlot()
    def send_message_all(self, client_msg: ClientMsgModel):
        """
        发送消息
        :param client_msg: 客户端消息
        :return:
        """
        for client_socket in self.client_socket_list:
            client_socket.write(client_msg.msg.encode())

    def close_all(self):
        """
        关闭所有连接
        :return:
        """
        for client_socket in self.client_socket_list:
            client_socket.close()
        self.close()

    @pyqtSlot()
    def close_client(self, client_msg: ClientMsgModel):
        """
        关闭指定连接
        :return:
        """
        for client_socket in self.client_socket_list:
            if (
                    client_socket.peerAddress().toString() == client_msg.client_ip
                    and client_socket.peerPort() == client_msg.client_port
            ):
                client_socket.close()

    @pyqtSlot()
    def close_client_all(self, client_msg: ClientMsgModel):
        """
        关闭指定连接
        :return:
        """
        for client_socket in self.client_socket_list:
            if (
                    client_socket.peerAddress().toString() == client_msg.client_ip
                    and client_socket.peerPort() == client_msg.client_port
            ):
                client_socket.close()


# TODO  UDP服务器
# TODO TCP客户端
# TODO UDP客户端


def main():
    """
    主函数，用于运行程序
    :return: None
    """
    app = QApplication(sys.argv)
    # st = StopThreading()
    ui = UiTcpUdp()
    ui.show()
    sys.exit(app.exec_())


def test_tcp_server():
    """
    测试Tcp服务器
    @return:
    """
    app = QApplication(sys.argv)

    server = TcpServer()
    server.signals_connected.connect(lambda x: print("connected", x))
    server.signals_disconnected.connect(lambda x: print("disconnected", x))
    server.signals_message.connect(lambda x: print("message", x))

    server.listen(QHostAddress("127.0.0.1"), 8888)
    print("server started")
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()
    # test_tcp_server()
